clc;close all;

dataPath = '/Users/MGHu/Documents/KRb reaction/Data/Figure 4 data/data/20190513_712nm.dat';
UVshotSig = 1:3500;
UVshotBkg = 3501:7000;

fnames = dir(dataPath);
disp('                              ');
disp(['---------',fnames.name,'--------']);
%% Set flags
massFlag = 1;
%%-----MCP Parameters--------------
TMCPacquire = 1*10^5; %set in the MCP acquisition software, in [ns]
dtMCPres = 0.025117; %TDC resolution [ns]
MCPRadius = 42;     %MCP detector radius [mm]
%%-----Parameters for data analysis--------
Xc = -10.250*0;      %[mm] center for plotting in detector frame
Yc = 8.75*0;    %[mm] center for plotting in detector frame
Angle = -30*0;   %[deg] rotation angle for plotting
tbinsize = 2; %5 20 [ns] ignore this if you are histcounts mass
Nbin = floor(TMCPacquire/tbinsize);
tEdge = (0:1:Nbin)*tbinsize; %[ns] for TOF binning
tList = (0.5:1:(Nbin-0.5))*tbinsize/1000; %[us]
% f(us) = 0.0214+ 2.481*sqrt(mass(amu));  %This formula from Yu's calibration
% massList = (tList/1000./15.69).^2*40-0.26;  %[amu] mass
if massFlag
    dmass = 0.25;
    massEdge = (20-dmass/2):dmass:(300+dmass/2);    %[amu]
    massList = 20:dmass:300;        %[amu]
    %     tList = 21.4 + 2481*sqrt(massList);  %[ns]
end
dXY = 0.15;        %[mm] spatial resolution
Xmin = -20;     %-45[mm] plot range
Ymin = Xmin+20;    %[mm]
Xmax = 0;      %45[mm]
Ymax = Xmax+20;    %[mm]
Xedge = Xmin:dXY:Xmax;
Yedge = Ymin:dXY:Ymax;
Xplot = (Xmin+dXY/2):dXY:(Xmax-dXY/2);
Yplot = (Ymin+dXY/2):dXY:(Ymax-dXY/2);
LabelFontSize = 15;
dmTOF = 6;      %% mass range for removing dominant peak in mass spectrum
mlist1 = [40 87 127];
species1 = {'K^+' 'Rb^+' 'KRb^+'};
dmPlot = 6;     %% mass range for plotting [amu]
mlistPlot = [127 87 40 80 174 214 167 254];%for spatial plot
speciesPlot = {'KRb^+' 'Rb^+' 'K^+' 'K_2^+' 'Rb_2^+' 'KRb_2^+' 'K_2Rb^+' 'K_2Rb_2^+'};%for spatial plot
speciesPlot2 = {'KRb' 'Rb' 'K' 'K2' 'Rb2' 'KRb2' 'K2Rb' 'K2Rb2'};
colorbg = zeros(1,3);       %Black
% colorbg = ones(1,3);        %white

Nfile = 1;
realHitList = cell(1, Nfile);
NUVpulse = zeros(1,Nfile);
dataTOFList = cell(1, Nfile);
dataMassList = cell(1, Nfile);
dataXYAllList = cell(1, Nfile);
dataXYPlotList = cell(1, Nfile);
datak = [];
for i = 1:Nfile
    filename = dataPath;
    if (~exist(filename,'file'))
        error(['Did not find such file!']);
    else
        disp('loading data ...');
        tic
        realHitList{i} = dlmread(filename,'\t');
        %1st column ---- TOF [ns]
        %2nd column ---- X position [mm]
        %3rd column ---- Y position [mm]
        %4th column ---- MCP trigger #
        %5th column ---- Exp cycle #
        NUVpulse(i) = max(UVshotBkg);
        dataTOF = zeros(NUVpulse(i),length(tList));
        dataMass = zeros(NUVpulse(i), length(massList));
        xtAll = realHitList{i}(:,1);
        XAll = realHitList{i}(:,2) - Xc;
        YAll = realHitList{i}(:,3) - Yc;
        dataXYAll = cell(1, NUVpulse(i));
        dataXYi = cell(1, NUVpulse(i));
        for j = 1:NUVpulse(i)
            xt = xtAll(realHitList{i}(:,4) == j); %[ns]
            dataTOF(j,:)= histcounts(xt, tEdge);
            xmass = ((xt-21.4)/2481).^2;  %[amu]
            dataMass(j,:) = histcounts(xmass, massEdge);
            %--All species at UVpulse j----
            dataX0 = XAll(realHitList{i}(:,4) == j);
            dataY0 = YAll(realHitList{i}(:,4) == j);
            dataX = dataX0*cos(-Angle*pi/180)-dataY0*sin(-Angle*pi/180);
            dataY = dataX0*sin(-Angle*pi/180)+dataY0*cos(-Angle*pi/180);
            dataXYAll{j} = histcounts2(dataY,dataX,Yedge,Xedge);
            dataXYj = zeros(length(Yplot),length(Xplot),length(mlistPlot));
            %--selected species at UVpulse j--------
            for k = 1:length(mlistPlot)
                dataRbX = dataX((xmass>=(mlistPlot(k)-dmPlot/2))&(xmass<=(mlistPlot(k)+dmPlot/2)));
                dataRbY = dataY((xmass>=(mlistPlot(k)-dmPlot/2))&(xmass<=(mlistPlot(k)+dmPlot/2)));
                dataXYj(:,:,k) = histcounts2(dataRbY,dataRbX,Yedge,Xedge);
            end
            dataXYi{j} = dataXYj;
        end
        dataTOFList{i} = dataTOF;       %[ns]
        dataXYAllList{i} = dataXYAll;      %[mm]
        dataXYPlotList{i} = dataXYi;
        dataMassList{i} = dataMass;     %[amu]
        toc
        %%%--- output some species data for other analysis---
        k = 6;
        disp(['--------Convert ', speciesPlot2{k} ,'.mat-----']);
        tic
        xt = realHitList{i}(:,1); %[ns]
        xmass = ((xt-21.4)/2481).^2;  %[amu]
        datak = [datak; realHitList{i}((xmass>=(mlistPlot(k)-dmPlot/2))&(xmass<=(mlistPlot(k)+dmPlot/2)),:)];
    end
end
dataKxRby = datak;
dataFileName = ['data/dataSpecies/', fnames.name(1:length(fnames.name)-4), '_', speciesPlot2{k},'.mat'];
save(dataFileName,'dataKxRby');
toc
close all;
%% figures
h(1) = figure('units','normalized','WindowStyle','docked','outerposition',[0 0 1 1]);
h(2) = figure('units','normalized','WindowStyle','docked','outerposition',[0 0 1 1]);
h(3) = figure('units','normalized','WindowStyle','docked','outerposition',[0 0 1 1]);
h(4) = figure('units','normalized','WindowStyle','docked','outerposition',[0 0 1 1]);
h(5) = figure('units','normalized','WindowStyle','docked','outerposition',[0 0 1 1]);
h(6) = figure('units','normalized','WindowStyle','docked','outerposition',[0 0 1 1]);

%%----Plot total TOF spectrometer-----------------
if 1
    figure(h(1));
    dataTOFSum = zeros(size(dataTOFList{1}(1,:)));
    for i=1:Nfile
        for j = UVshotSig
            dataTOFSum = dataTOFSum+dataTOFList{i}(j,:);
        end
    end
    mTOF = tList;
    countsTOF1 = dataTOFSum;
    subplot(2,1,1)
    plot(tList, countsTOF1,'Color','black');
    xlim([11,35]);
    xlabel('TOF (us)');
    ylabel('Ion (count)');
    title(['TOF spectrum of ions']);
    xt = get(gca, 'XTick');
    set(gca, 'FontSize', LabelFontSize);
    xt = get(gca, 'YTick');
    set(gca, 'FontSize', LabelFontSize);
    subplot(2,1,2)
    dataMassSum = zeros(size(dataMassList{1}(1,:)));
    for i = 1:Nfile
        for j = UVshotSig
            dataMassSum = dataMassSum+dataMassList{i}(j,:);
        end
    end
    countsTOF1 = dataMassSum;
    plot(massList, countsTOF1,'Color','black');
    xlim([20,300]);
    xlabel('m/Z');
    ylabel('Ion (count)');
    %     title(['Mass spectrum of ions']);
    xt = get(gca, 'XTick');
    set(gca, 'FontSize', LabelFontSize);
    xt = get(gca, 'YTick');
    set(gca, 'FontSize', LabelFontSize);
end


%%----Plot total Mass spectrometer with and without KRb -----------------
if 1
    figure(h(2));
    dataMassSum = zeros(size(dataMassList{1}(1,:)));
    subplot(2,1,1)
    for i=1:Nfile
        for j=UVshotSig
            dataMassSum = dataMassSum+dataMassList{i}(j,:);
        end
    end
    mTOF = massList;
    countsTOF1 = dataMassSum;
    for i = 1:length(mlist1)
        countsTOF1(mTOF>=(mlist1(i)-dmTOF/2)& mTOF<=(mlist1(i)+dmTOF/2)) = 0;
    end
    plot(massList, countsTOF1,'Color','black');
    hold on
    xlim([20,300]);
    mlist = [40 87 127];
    species = {'K^+', 'Rb^+', 'KRb^+'};
    disp('Ion signal:');
    for i = 1:length(mlist)
        xdata = mTOF(mTOF>=(mlist(i)-dmTOF/2)&mTOF<=(mlist(i)+dmTOF/2));
        ydata = dataMassSum(mTOF>=(mlist(i)-dmTOF/2)&mTOF<=(mlist(i)+dmTOF/2));
        plot(xdata, ydata,'-','LineWidth',1, 'Color', [0 0 0]+0.8);
        hold on;
        disp(['  N', species{i},'=',num2str(sum(ydata))]);
    end
    %     dmTOF = 11;
    mlist = [80 174 167 214 254];
    species = {'K_2^+' 'Rb_2^+' 'K_2Rb^+' 'KRb_2^+' 'K_2Rb_2^+'};
    colorlist = {'red' 'cyan' 'green' 'magenta' 'blue'};
    yplotmax = 0;
    for i = 1:length(mlist)
        xdata = mTOF(mTOF>=(mlist(i)-dmTOF/2)&mTOF<=(mlist(i)+dmTOF/2));
        ydata = dataMassSum(mTOF>=(mlist(i)-dmTOF/2)&mTOF<=(mlist(i)+dmTOF/2));
        plot(xdata, ydata, '-', 'LineWidth', 0.75, 'Color', colorlist{i});
        if strcmp(species{i},'K_2Rb^+')
            text(mlist(i)-22, 5+max(ydata),species{i}, 'Color',colorlist{i});
        elseif strcmp(species{i},'K_2^+')
            text(mlist(i)-10, 5+max(ydata),species{i}, 'Color',colorlist{i});
        else
            text(mlist(i)-10, 5+max(ydata),species{i}, 'Color',colorlist{i});
        end
        yplotmax = max(yplotmax, max(ydata));
        hold on;
        disp(['  N', species{i},'=',num2str(sum(ydata))]);
    end
    ylim([0, yplotmax*1.5]);
    hold off
    %     title(['UV pulse time after STIRAP: ', num2str((UVpulseNoList(k)-1)*100+15), ' ms']);
    xlabel('m/Z');
    ylabel('Ion (count)');
    xt = get(gca, 'XTick');
    set(gca, 'FontSize', LabelFontSize);
    xt = get(gca, 'YTick');
    set(gca, 'FontSize', LabelFontSize);
    
    subplot(2,1,2)
    dataMassSum = zeros(size(dataMassList{1}(1,:)));
    for i = 1 : Nfile
        for j = UVshotBkg
            dataMassSum = dataMassSum+dataMassList{i}(j,:);
        end
    end
    countsTOF1 = dataMassSum;
    for i = 1:length(mlist1)
        countsTOF1(mTOF>=(mlist1(i)-dmTOF/2)& mTOF<=(mlist1(i)+dmTOF/2)) = 0;
    end
    plot(massList, countsTOF1,'Color','black');
    hold on
    xlim([20, 300]);
    ylim([0, yplotmax*1.5]);
    mlist = [40 87 127];
    species = {'K^+', 'Rb^+', 'KRb^+'};
    disp('Ion background:');
    for i = 1:length(mlist)
        xdata = mTOF(mTOF>=(mlist(i)-dmTOF/2)&mTOF<=(mlist(i)+dmTOF/2));
        ydata = dataMassSum(mTOF>=(mlist(i)-dmTOF/2)&mTOF<=(mlist(i)+dmTOF/2));
        plot(xdata, ydata,'-','LineWidth',1, 'Color', [0 0 0]+0.8);
        hold on;
        disp(['  N', species{i},'=',num2str(sum(ydata))]);
    end
    mlist = [80 174 167 214 254];
    species = {'K_2^+' 'Rb_2^+' 'K_2Rb^+' 'KRb_2^+' 'K_2Rb_2^+'};
    colorlist = {'red' 'cyan' 'green' 'magenta' 'blue'};
    for i = 1:length(mlist)
        xdata = mTOF(mTOF>=(mlist(i)-dmTOF/2)&mTOF<=(mlist(i)+dmTOF/2));
        ydata = dataMassSum(mTOF>=(mlist(i)-dmTOF/2)&mTOF<=(mlist(i)+dmTOF/2));
        plot(xdata, ydata,'-','LineWidth', 1.0, 'Color', colorlist{i});
        hold on;
        disp(['  N', species{i},'=',num2str(sum(ydata))]);
    end
    hold off
    xlabel('m/Z');
    ylabel('Ion (count)');
    xt = get(gca, 'XTick');
    set(gca, 'FontSize', LabelFontSize);
    xt = get(gca, 'YTick');
    set(gca, 'FontSize', LabelFontSize);
end

%%-----Plot All XY distribution of selected species--------------
if 1
    figure(h(3));
    LabelFontSize = 10;
    for i = 1:length(mlistPlot)
        dataXYPlot = zeros(length(Yplot),length(Xplot));
        for j = 1:Nfile
            for k = UVshotSig
                dataXYPlot = dataXYPlot+dataXYPlotList{j}{k}(:,:,i);
            end
        end
        subplot(3,3,i)
        maxCounts = ceil(max(max(dataXYPlot)));
        dataXYPlotCut = dataXYPlot;
        %         if i==2
        %             dataXYPlotCut(dataXYPlotCut > maxCounts/160) = maxCounts/160;
        %         end
        %         if i==3
        %             dataXYPlotCut(dataXYPlotCut > maxCounts/30) = maxCounts/30;
        %         end
        if maxCounts < 2
            dataXYPlotCut(1,1) = 2;
        end
        imagesc(Xplot, Yplot, dataXYPlotCut);
        if i == 1
            cmap = jet(maxCounts);
        end
        cmap(1:1,:) = colorbg;    % Make values 0 black:
        colormap(cmap);
        colorbar
        xlabel('x (mm)');
        ylabel('y (mm)');
        title(['All ',speciesPlot{i},': m=',num2str(mlistPlot(i)), '\pm', num2str(dmPlot/2), ' amu']);
        set(gca,'YDir','normal');
        xtick = get(gca, 'XTick');
        set(gca, 'FontSize', LabelFontSize);
        ytick = get(gca, 'YTick');
        set(gca, 'FontSize', LabelFontSize);
        hold on
        t = linspace(0,2*pi);
        plot(MCPRadius.*cos(t),MCPRadius.*sin(t),'w--');
    end
end

%%-----plot selected species--------
close(h(4));
h(4) = figure('units','normalized','WindowStyle','docked','outerposition',[0 0 1 1]);
figure(h(4));
LabelFontSize = 15;
i = 8;
dataXYPlot = zeros(length(Yplot),length(Xplot));
for j = 1:Nfile
    for k = UVshotSig
        dataXYPlot = dataXYPlot+dataXYPlotList{j}{k}(:,:,i);
    end
end
maxCounts = ceil(max(max(dataXYPlot)));
dataXYPlotCut = dataXYPlot;
if i==2
    dataXYPlotCut(dataXYPlotCut > maxCounts/200) = maxCounts/200;
end
if i==3
    dataXYPlotCut(dataXYPlotCut > maxCounts/20) = maxCounts/20;
end
if i==8
    dataXYPlotCut(1,1) = 25;
end
if maxCounts < 2
    dataXYPlotCut(1,1) = 2;
end
imagesc(Xplot, Yplot, dataXYPlotCut);
cmap(1:1,:) = colorbg;    % Make values 0 black:
colormap(cmap);
colorbar
xlabel('x (mm)');
ylabel('y (mm)');
% title([speciesPlot{i},': m=',num2str(mlistPlot(i)), '\pm', num2str(dmPlot/2), ' amu']);
set(gca,'YDir','normal');
xtick = get(gca, 'XTick');
set(gca, 'FontSize', LabelFontSize);
ytick = get(gca, 'YTick');
set(gca, 'FontSize', LabelFontSize);
hold on;

xplc = -10.93;
yplc = 9.525;
dx = 8;
xlim([xplc-dx/2,xplc+dx/2]);
ylim([yplc-dx/2,yplc+dx/2]);

t = linspace(0,2*pi);
KE = 10.4;       %  cm^-1
VR = 1000;          % [V] repeller voltage
AVMI = 16.44;        %[mm/sqrt(cm-1/V)]
R10cm = AVMI*sqrt(KE/VR);
plot(R10cm.*cos(t)+xplc, R10cm.*sin(t)+yplc, 'y--','LineWidth', 1.5);
hold off

%%-----Plot TOF lineshape--------------
if 1
    LabelFontSize = 10;
    figure(h(5));
    dataTOFSum = zeros(size(dataTOFList{1}(1,:)));
    for k= 1:Nfile
        for j = UVshotSig
            dataTOFSum = dataTOFSum+dataTOFList{k}(j,:);
        end
    end
    for i = 1:length(mlistPlot)
        dTOF =0.5; % [us]
        xTOF = (sqrt(mlistPlot(i))*2481+21.4)/1000;   %[us]
        %             xmass = ((xt-21.4)/2481).^2;  %[amu]
        xdata = tList(tList>=(xTOF-dTOF/2)&tList<=(xTOF+dTOF/2));%[us]
        ydata = dataTOFSum(tList>=(xTOF-dTOF/2)&tList<=(xTOF+dTOF/2));
        subplot(3, 3, i)
        plot(xdata, ydata,'o-','LineWidth', 1.0, 'MarkerSize', 5);
        title([speciesPlot{i}]);
        xlim([xTOF-dTOF/2, xTOF+dTOF/2]);
        xlabel('TOF (\mus)');
        ylabel('Ion (count)');
        xticks([xTOF-dTOF/2, xTOF, xTOF+dTOF/2])
        xticklabels([xTOF-dTOF/2, xTOF, xTOF+dTOF/2])
        xt = get(gca, 'XTick');
        set(gca, 'FontSize', LabelFontSize);
        xt = get(gca, 'YTick');
        set(gca, 'FontSize', LabelFontSize);
    end
end

%%-----Plot combined density-dependent mass spectrum-----------------------
if 1
    figure(h(6));
    LabelFontSize = 15;
    mlist1 = [127 167 214 254];
    species1 = {'KRb^+' 'K_2Rb^+' 'KRb_2^+' 'K_2Rb_2^+'};
    frep = 7000;    %[Hz]
    dt = 1/frep*1000;    %[s]
    Dt = 15*dt;
    tUV = UVshotSig*dt;
    tList = (1:floor(max(UVshotSig)*dt/Dt))*Dt; %[s]
    dmTOF = 6;
    mTOF = massList;
    for i = 1:length(mlist1)
        massUV = zeros(size(UVshotSig));
        for UVpulseNo = UVshotSig
            dataMassSum = dataMassList{1}(UVpulseNo,:);
            massUV(UVpulseNo) = sum(dataMassSum(mTOF>=(mlist1(i)-dmTOF/2) & mTOF<=(mlist1(i)+dmTOF/2)));
        end
        massUV1 = zeros(size(tList));
        for j = 1:floor(max(UVshotSig)*dt/Dt)
            massUV1(j) = sum(massUV(tUV >= tList(j)-Dt & tUV < tList(j)));
        end
        if i == 1
            massKRb = massUV1;
        end
        if i == 3
            massKRb2 = massUV1;
        end
        if i == 4
            massK2Rb2 = massUV1;
        end
        subplot(3,2,i)
        plot(tList, massUV1, '-');
        xlim([1 400]);
        title(species1{i});
        xlabel('time (ms)');
        ylabel('Ion (count)');
        xt = get(gca, 'XTick');
        set(gca, 'FontSize', LabelFontSize);
        xt = get(gca, 'YTick');
        set(gca, 'FontSize', LabelFontSize);
    end
    subplot(3,2,[5 6])
    NKRb2norm = massKRb2./massKRb.*massKRb(1);
    NK2Rb2norm = massK2Rb2./massKRb.*massKRb(1);
    plot(tList, NKRb2norm, '-m');
    hold on
    plot(tList, NK2Rb2norm, '-b');
    xmin = 1;
    xmax = 80;
    xlim([xmin xmax]);
    NtotKRb2 = sum(NKRb2norm(tList >= xmin & tList <= xmax));
    NtotK2Rb2 = sum(NK2Rb2norm(tList >= xmin & tList <= xmax));
    title(['KRb_2^+/KRb^+ = ', num2str(NtotKRb2), ', K_2Rb_2^+/KRb^+ = ', num2str(NtotK2Rb2)]);
    disp(['KRb_2^+/KRb^+ = ', num2str(NtotKRb2), ', K_2Rb_2^+/KRb^+ = ', num2str(NtotK2Rb2)]);
    xlabel('time (ms)');
    ylabel('Ion (count)');
    hold off
end

figure(h(4));
% saveas(h(4), ['plots/VMI_K2Rb2_356nm.png']);
% savefig(h, ['plots/', fnames.name(1:length(fnames.name)-4), '.fig']);
% saveas(h(2), ['plots/', fnames.name(1:length(fnames.name)-4), '_mass_spectrum.png']);
% saveas(h(3), ['plots/', fnames.name(1:length(fnames.name)-4), '_VMI_All.png']);
% saveas(h(4), ['plots/', fnames.name(1:length(fnames.name)-4), '_VMI_K2Rb2.png']);
% saveas(h(5), ['plots/', fnames.name(1:length(fnames.name)-4), '_TOF_lineshape.png']);


